# Run entire code at once!
# Legal Aid and Wrongful Conviction (frontdoor)
setwd("/Users/firatbilgel/Documents/Journal articles/Identification")
libraries <- c("mvtnorm", "pbivnorm", "endogeneity", "boot", "haven", "MASS", "dplyr", "kableExtra", "modelsummary", "knitr")
for (library_name in libraries) {
  library(library_name, character.only = TRUE)
}
rm(list = ls())
set.seed(123)
n <- 10000
ep1 <- rnorm(n, mean = 0, sd = 1)
ep2 <- rnorm(n, mean = 0, sd = 1)
ep3 <- rnorm(n, mean = 0, sd = 1)
ep4 <- rnorm(n, mean = 0, sd = 1)
ep5 <- rnorm(n, mean = 0, sd = 1)
ep6 <- rnorm(n, mean = 0, sd = 1)
ep7 <- rnorm(n, mean = 0, sd = 1)
U1 <- pmin(pmax(rpois(n, lambda = 12), 0), 23) # family SES (schooling attainment, years)
U2 <- pmin(pmax(rpois(n, lambda = 4), 1), 10) # community wealth (quintiles)
U3 <- pmin(pmax(rpois(n, lambda = 3), 1), 5) # legal system bias
U4 <- pmin(pmax(rpois(n, lambda = 3), 1), 5) # judicial discretion
Z1 <- pmin(pmax(rpois(n, lambda = exp(0.5*U2 + ep1)), 1), 5) # knowledge of the legal system (likert-type)
Z2 <- pmin(pmax(rpois(n, lambda = exp(0.5*Z1 + ep2)), 1), 5) # quality of legal representation (likert-type)
X1 <- pmin(pmax(rpois(n, lambda = exp(-0.4*Z2 - 0.5*U1 + 0.6*U4 + ep3)), 0), 5) # arrest circumstances
X2 <- pmin(pmax(rpois(n, lambda = exp(-0.5*U1 - 0.5*U2 + 2*U3 + ep4)), 0), 40) # criminal history
alpha = -0.2
delta = -0.65
gamma = 5.5
beta = -3.6
# Treatment (fd-iv)
lincomT <- alpha - 0.2 * X1 + 0.45 * X2 + 0.15 * Z1 + 0.05 * U1 + 0.05 * U2 + 0.3 * U3 + ep5 
probT <- pmin(pmax(lincomT, 0), 1) 
Treatment <- rbinom(n, 1, probT) # access to legal aid (binary) 
# Mediator (fd)
Q <- pmin(pmax(rpois(n, lambda = exp(gamma*Treatment + ep7)), 1), 5) # rigour of defense strategy
# Outcome (fd)
lincomY <- delta * Q + 0.3 * X1 + 0.5 * U3 + 0.03 * U4 + ep6 
probY <- pmin(pmax(lincomY, 0), 1)
Y <- rbinom(n, 1, probY) # wrongful conviction (binary) 

betamax2=-3.6
# Outcome (iv)
lincomY_iv <- betamax2* Treatment + 0.3 * X1 + 0.5 * U3 + 0.03 * U4 + ep6
probY_iv <- pmin(pmax(lincomY_iv, 0), 1)
Y_iv <- rbinom(n, 1, probY_iv)
dffrontdoor <- data.frame(Z1, Z2, X1, X2, Treatment, Y_iv, Y, Q, U1, U2, U3, U4)
dffrontdoor$diff <- Y_iv-Y
summary(dffrontdoor)
summary <- dffrontdoor[, c("Z1","Z2", "X1", "X2", "Treatment", "Y_iv", "Y", "Q", "U1", "U2", "U3", "U4")] %>% 
  select(`Knowledge of the legal system` = Z1, `Quality of legal representation` = Z2, `Arrest circumstances` = X1, `Criminal history` = X2, 
         `Access to legal aid` = Treatment, `Wrongful conviction IV` = Y_iv, `Wrongful conviction FD` = Y, `Rigor of defense strategy` = Q, `family SES` = U1, `Community wealth` = U2, `Legal system bias` = U3, `Judicial discretion` = U4)
summary_list <- lapply (summary, scale)
emptycol = function(x) " "
datasummary(`Knowledge of the legal system` + `Quality of legal representation` + `Arrest circumstances` +`Criminal history` +
              `Access to legal aid` + `Wrongful conviction IV` + `Wrongful conviction FD` + `Rigor of defense strategy` + `family SES` + `Community wealth` + `Legal system bias` + `Judicial discretion`~  Mean + SD + Min + Max + Heading("Boxplot") * emptycol + Heading("Histogram") * emptycol, data=summary, fmt = "%.3f") %>% column_spec(column = 6, image = spec_boxplot(summary_list)) %>% column_spec(column = 7, image =spec_hist(summary_list))
write_dta(dffrontdoor, "lawc_impossible.dta")

# Define potential outcomes under treatment and no treatment (iv)
probY_treated_iv <- pmin(pmax(betamax2 * 1 + 0.3 * X1 + 0.5 * U3 + 0.03 * U4 + ep6, 0), 1)
Y_treated_iv <- rbinom(n, 1, probY_treated_iv)
probY_control_iv <- pmin(pmax(betamax2 * 0 + 0.3 * X1 + 0.5 * U3 + 0.03 * U4 + ep6, 0), 1)
Y_control_iv <- rbinom(n, 1, probY_control_iv)
true_ATT_iv <- mean(Y_treated_iv[Treatment == 1] - Y_control_iv[Treatment == 1])

# Define potential outcomes under treatment and no treatment (fd)
Q_Treated <- pmin(pmax(rpois(n, lambda = exp(gamma * 1 + ep7)), 1), 5)
lincomY_Treated <- delta * Q_Treated + 0.3 * X1 + 0.5 * U3 + 0.03 * U4 + ep6
probY_Treated <- pmin(pmax(lincomY_Treated, 0), 1)
Y_Treated <- rbinom(n, 1, probY_Treated)

Q_Control <- pmin(pmax(rpois(n, lambda = exp(gamma * 0 + ep7)), 1), 5)
lincomY_Control <- delta * Q_Control + 0.3 * X1 + 0.5 * U3 + 0.03 * U4 + ep6
probY_Control <- pmin(pmax(lincomY_Control, 0), 1)
Y_Control <- rbinom(n, 1, probY_Control)
true_ATT_fd <- mean(Y_Treated[Treatment == 1] - Y_Control[Treatment == 1])

######### Biprobit
invisible(capture.output(fit <- biprobit(Treatment ~ X1 + Z2 + Z1, Y_iv ~ Treatment + X1 + Z2, data = dffrontdoor, method = "BFGS")))
coef_selection <- fit$par[grep("^1", names(fit$par))]
coef_outcome <- fit$par[c("(Intercept)", "Treatment", "X1", "Z2")]
tau <- fit$par["tau"]

# Predict latent variables for selection equation
xb1 <- as.vector(model.matrix(~X1 + Z2 + Z1, data = dffrontdoor) %*% coef_selection)
# Predict latent variables for outcome equation
xb2 <- as.vector(model.matrix(~Treatment + X1 + Z2, data = dffrontdoor) %*% coef_outcome)
# Compute ATT for treated individuals
tev <- (pbivnorm(cbind(xb1, xb2), rho = tau) - pbivnorm(cbind(xb1, xb2 - coef_outcome["Treatment"]), rho = tau)) /
  pnorm(xb1)
# Filter for treated individuals
tev_treated <- tev[dffrontdoor$Treatment == 1]
# Calculate ATT
iv_ATT <- mean(tev_treated, na.rm = TRUE)
ATT_bias_iv <- (iv_ATT-true_ATT_iv)/true_ATT_iv*100

######### Biprobit (invalid estimates due to includion of X2)
invisible(capture.output(fit_invalid <- biprobit(Treatment ~ X1 + Z2 + Z1 + X2, Y_iv ~ Treatment + X1 + Z2 + X2, data = dffrontdoor, method = "BFGS")))
coef_selection_invalid <- fit$par[grep("^1", names(fit_invalid$par))]
coef_outcome_invalid <- fit_invalid$par[c("(Intercept)", "Treatment", "X1", "Z2", "X2")]
tau_invalid <- fit_invalid$par["tau"]
# Constrain tau to valid range
tau_invalid <- max(-1, min(1, tau_invalid))
# Predict latent variables for selection equation
xb1_invalid <- as.vector(model.matrix(~X1 + Z2 + Z1 + X2, data = dffrontdoor) %*% coef_selection_invalid)
# Predict latent variables for outcome equation
xb2_invalid <- as.vector(model.matrix(~Treatment + X1 + Z2 + X2, data = dffrontdoor) %*% coef_outcome_invalid)
# Compute ATT for treated individuals
tev_invalid <- (pbivnorm(cbind(xb1_invalid, xb2_invalid), rho = tau_invalid) - pbivnorm(cbind(xb1_invalid, xb2_invalid - coef_outcome_invalid["Treatment"]), rho = tau_invalid)) /
  pnorm(xb1_invalid)
# Filter for treated individuals
tev_treated_invalid <- tev_invalid[dffrontdoor$Treatment == 1]
# Calculate ATT
iv_ATT_invalid <- mean(tev_treated_invalid, na.rm = TRUE)
ATT_bias_iv_invalid <- (iv_ATT_invalid-true_ATT_iv)/true_ATT_iv*100

########## Step 1: Ordered logit, Treatment on Q
step1 <- polr(as.factor(Q) ~ Treatment, data = dffrontdoor, method = "probit")
summary(step1)
########## Step 2: Logit, Q on Y, controlling for Treatment
step2 <- glm(Y ~ Q + Treatment, family = binomial(link = "logit"), data = dffrontdoor)
summary(step2)
# Filter for treated individuals before predicting Q_Treated
Q_Treated_filtered <- predict(step1, newdata = dffrontdoor %>% filter(Treatment == 1), type = "probs")
Q_Treated_expected_filtered <- apply(Q_Treated_filtered, 1, function(p) sum(p * 1:5)) # Expected value of Q for treated
# Filter for treated individuals before predicting Y_Treated
Y_Treated_ATT <- predict(step2, newdata = dffrontdoor %>% filter(Treatment == 1) %>% mutate(Q = Q_Treated_expected_filtered), type = "response")
# Predict Y_Control for treated individuals (what would happen if they had not been treated)
Q_Control_filtered <- predict(step1, newdata = dffrontdoor %>% filter(Treatment == 1) %>% mutate(Treatment = 0), type = "probs")
Q_Control_expected_filtered <- apply(Q_Control_filtered, 1, function(p) sum(p * 1:5))
Y_Control_ATT <- predict(step2, newdata = dffrontdoor %>% filter(Treatment == 1) %>% mutate(Q = Q_Control_expected_filtered), type = "response")

# Calculate ATT
fd_ATT <- mean(Y_Treated_ATT - Y_Control_ATT)
ATT_bias_fd <- (fd_ATT-true_ATT_fd)/true_ATT_fd*100

cat("True ATT iv: ", true_ATT_iv, "\n")
cat("IV ATT:", iv_ATT, "\n")
cat("IV ATT invalid:", iv_ATT_invalid, "\n")
cat("ATT Bias (iv):", ATT_bias_iv, "\n")
cat("ATT Bias (invalid iv):", ATT_bias_iv_invalid, "\n")
cat("True ATT fd: ", true_ATT_fd, "\n")
cat("Frontdoor ATT: ", fd_ATT, "\n")
cat("ATT Bias (Frontdoor): ", ATT_bias_fd, "\n")

##################################################
# Define a Function to compute ATT via Frontdoor #
##################################################
att_boot <- function(data, indices) {
  d <- data[indices,]  # Resample the data
  # Re-run the models on bootstrapped data
  step1 <- polr(as.factor(Q) ~ Treatment, data = d, method = "probit")
  step2 <- glm(Y ~ Q + Treatment, family = binomial(link = "logit"), data = d)
  # Predicted values for Q given Treatment = 1 and Treatment = 0
  Q_Treated_filtered <- predict(step1, newdata = d %>% filter(Treatment == 1), type = "probs")
  Q_Control_filtered <- predict(step1, newdata = d %>% filter(Treatment == 1) %>% mutate(Treatment = 0), type = "probs")
  # Calculate expected value of Q for treated and control
  Q_Treated_expected_filtered <- apply(Q_Treated_filtered, 1, function(p) sum(p * 1:5))
  Q_Control_expected_filtered <- apply(Q_Control_filtered, 1, function(p) sum(p * 1:5))
  # Predicted probabilities for Y for treated individuals
  Y_Treated_ATT <- predict(step2, newdata = d %>% filter(Treatment == 1) %>% mutate(Q = Q_Treated_expected_filtered), type = "response")
  Y_Control_ATT <- predict(step2, newdata = d %>% filter(Treatment == 1) %>% mutate(Q = Q_Control_expected_filtered), type = "response")
  # Calculate ATT
  mean(Y_Treated_ATT - Y_Control_ATT)
}

# Bootstrap for ATT (fd)
set.seed(123)
n_boot <- 1000
boot_att_fd <- boot(data = dffrontdoor, statistic = att_boot, R = n_boot)
# Confidence intervals for ATT
ci_att_fd <- boot.ci(boot_att_fd, type = c("perc"))
print(ci_att_fd)

################################################
# Define a function to compute ATT via Biprobit
################################################
compute_att <- function(data, indices) {
  # Resample the data
  d <- data[indices, ]
  # Fit the biprobit model
  invisible(capture.output(fit <- biprobit(Treatment ~ X1 + Z2 + Z1, Y_iv ~ Treatment + X1 + Z2, data = d, method = "BFGS")))
  # Extract coefficients
  coef_selection <- fit$par[grep("^1", names(fit$par))]
  coef_outcome <- fit$par[c("(Intercept)", "Treatment", "X1", "Z2")]
  tau <- fit$par["tau"]
  # Constrain tau to valid range
  tau <- max(-1, min(1, tau))
  # Predict latent variables for selection and outcome equations
  xb1 <- as.vector(model.matrix(~X1 + Z2 + Z1, data = d) %*% coef_selection)
  xb2 <- as.vector(model.matrix(~Treatment + X1 + Z2, data = d) %*% coef_outcome)
  # Compute ATT for treated individuals
  tev <- (pbivnorm(cbind(xb1, xb2), rho = tau) - pbivnorm(cbind(xb1, xb2 - coef_outcome["Treatment"]), rho = tau)) / pnorm(xb1)
  # Filter for treated individuals
  tev_treated <- tev[d$Treatment == 1]
  # Return ATT
  return(mean(tev_treated, na.rm = TRUE))
}

# Bootstrap for ATT (iv)
set.seed(123)  
n_boot <- 1000 
boot_att_iv <- boot(data = dffrontdoor, statistic = compute_att, R = n_boot)
# Confidence intervals
ci_att_iv <- boot.ci(boot_att_iv, type = c("perc"))
print(ci_att_iv)

#####################################################################
# Define a function to compute ATT via Biprobit (invalid estimates) #
#####################################################################
compute_att_invalid <- function(data, indices) {
  # Resample the data
  d <- data[indices, ]
  # Fit the biprobit model
  invisible(capture.output(fit_invalid <- biprobit(Treatment ~ X1 + Z2 + Z1 + X2, Y_iv ~ Treatment + X1 + Z2 + X2, data = d, method = "BFGS")))
  # Extract coefficients
  coef_selection_invalid <- fit$par[grep("^1", names(fit_invalid$par))]
  coef_outcome_invalid <- fit_invalid$par[c("(Intercept)", "Treatment", "X1", "Z2", "X2")]
  tau_invalid <- fit_invalid$par["tau"]
  # Constrain tau to valid range
  tau_invalid <- max(-1, min(1, tau_invalid))
  # Predict latent variables for selection and outcome equations
  xb1_invalid <- as.vector(model.matrix(~X1 + Z2 + Z1 + X2, data = d) %*% coef_selection_invalid)
  xb2_invalid <- as.vector(model.matrix(~Treatment + X1 + Z2 + X2, data = d) %*% coef_outcome_invalid)
  # Compute ATT for treated individuals
  tev_invalid <- (pbivnorm(cbind(xb1_invalid, xb2_invalid), rho = tau_invalid) - pbivnorm(cbind(xb1_invalid, xb2_invalid - coef_outcome_invalid["Treatment"]), rho = tau_invalid)) / pnorm(xb1_invalid)
  # Filter for treated individuals
  tev_treated_invalid <- tev_invalid[d$Treatment == 1]
  # Return ATT
  return(mean(tev_treated_invalid, na.rm = TRUE))
}

# Bootstrap for ATT (iv-invalid)
set.seed(123) 
n_boot <- 1000
boot_att_iv_invalid <- boot(data = dffrontdoor, statistic = compute_att_invalid, R = n_boot)
# Confidence intervals
ci_att_iv_invalid <- boot.ci(boot_att_iv_invalid, type = c("perc"))
print(ci_att_iv_invalid)
